import PyPDF2
import re
import os
from tqdm import tqdm

def extract_file_paths(pdf_path):
    print(f"正在处理PDF文件: {pdf_path}")
    # 打开PDF文件
    with open(pdf_path, 'rb') as file:
        # 创建PDF读取器对象
        pdf_reader = PyPDF2.PdfReader(file)
        total_pages = len(pdf_reader.pages)
        print(f"PDF文件页数: {total_pages}")
        
        # 用于存储找到的文件路径和行号
        findings = []  # 使用list存储所有匹配项
        
        # 使用tqdm创建进度条
        with tqdm(total=total_pages, desc="处理进度", unit="页") as pbar:
            # 遍历每一页
            for page_num, page in enumerate(pdf_reader.pages, 1):
                # 提取文本
                text = page.extract_text()
                
                # 使用正则表达式查找以"缺陷标识:"开头的行中的文件路径和行号
                # 处理中文PDF文档中的各种空白字符和换行
                matches = re.finditer(r'缺陷标识:[ \t\r\n\f\v]*(tmp/code/panweidb_server/(?:[^:\r\n]+/)*[^:\r\n]+):(\d+)(?:\s*\([^)]*\))?', text, re.MULTILINE)
                
                for match in matches:
                    file_path = match.group(1)
                    line_number = match.group(2)
                    findings.append(f"{file_path}:{line_number}")
                
                # 更新进度条
                pbar.update(1)
                # 显示当前找到的匹配数
                pbar.set_postfix({"已找到": len(findings)})
        
        return findings

def main():
    # PDF文件路径 - 使用绝对路径
    current_dir = os.path.dirname(os.path.abspath(__file__))
    pdf_path = os.path.join(current_dir, 'panweidb_server_security.pdf')
    output_file = os.path.join(current_dir, 'security_findings.txt')
    
    try:
        results = extract_file_paths(pdf_path)
        
        print("\n正在写入结果到文件...")
        # 将结果写入文件，只写入路径和行号
        with open(output_file, 'w', encoding='utf-8') as f:
            for path_with_line in results:
                f.write(f"{path_with_line}\n")
            
        print(f"结果已保存到文件: {output_file}")
        print(f"共找到 {len(results)} 个缺陷标识")
            
    except FileNotFoundError:
        print(f"错误: 找不到文件 {pdf_path}")
    except Exception as e:
        print(f"处理PDF时发生错误: {str(e)}")

if __name__ == '__main__':
    main() 